import sys, getopt
from PIL import Image
from pathlib import Path
import os
import json

# 支持的图片后缀
SUFFIX = ['.jpg', '.jpeg','.png']


def get_executable_dir():
    # 如果是通过PyInstaller打包后的exe运行
    if getattr(sys, 'frozen', False):
        # 获取exe文件的目录
        return os.path.dirname(sys.executable)
    else:
        # 正常Python脚本运行时
        return os.path.dirname(os.path.abspath(__file__))


def convert_mb_kb(bytesize):
    """
    把byte长度转换成KB,MB
    """
    if bytesize > 0:
        bytesize = bytesize / 1024
        if bytesize < 1024:
            return "%.fKB" % bytesize
        else:
            return "%.2fMB" % (bytesize / 1024) 

# exe 或脚本执行目录
current_dir = get_executable_dir()


class ImageCompression:
    def __init__(self, argv):
        self.images_info = []
        self._get_options(argv)
    
    def _help(self):
        '''命令提示
        '''
        print('\n')
        print('test.py [-i <imgs>] [-q <quality>] [-s <subsampling>] [-j <jpga>] [-d <dir>] [-o <output>]')
        print('     -i, --imgs 需要压缩的图片，多个图片以逗号分隔 "a.jpg,b.jpg')
        print('     -q, --quality 默认压缩的图片质量为15，可以调整0-95 ')
        print('     -j, --jpga 为1时设置将图片统计转换成.jpg格式，默认为0 ')
        print('     -d, --dir 设置一个目录，压缩指定目录下的图片 ')
        print('     -r, --recursion 是否递归遍历指定的文件夹, 默认否0 ')
        print('     -o, --output 设置一个目录，图片压缩后输出目录,默认图片所在目录的output文件夹 ')
        print('     -l, --log 是否展示执行过程, 默认0 ')
        print('     -s, subsampling 设置编码器的子采样 默认-1 ')
        print('                     -1: equivalent to keep ')
        print('                      0: equivalent to 4:4:4 ')
        print('                      1: equivalent to 4:2:2 ')
        print('                      2: equivalent to 4:2:0 ')
        print('\n')
        print('命令示例：python test.py -i a.jpg,b.jpg -q 20')
        print('\n')

    def _get_options(self, argv):
        """解析命令行参数"""
        try:
            self.imgs = ''
            self.quality = 15  # 图片质量
            self.jpga = 0  # 是否转为jpg
            self.log = 0  # 不展示执行过程
            self.recursion = 0  # 是否递归遍历文件夹
            self.subsampling = -1
            self.dir = ''  # 图片目录
            self.output = os.path.join(current_dir, 'output')  # 默认执行目录下的output文件夹

            opts, args = getopt.getopt(argv, "i:q:s:j:d:o:r:l:", ["imgs=", "quality=", "subsampling=", "jpga=", "dir=", "output=", "recursion=", "log="])
        except getopt.GetoptError:
            self._help()
            sys.exit(2)
        else:
            for opt, arg in opts:
                if opt == '-h':
                    self._help()
                    sys.exit()
                elif opt in ("-i", "--imgs"):
                    self.imgs = arg
                elif opt in ("-q", "--quality"):
                    self.quality = int(arg)
                elif opt in ("-s", "--subsampling"):
                    self.subsampling = int(arg)
                elif opt in ("-j", "--jpga"):
                    self.jpga = int(arg)
                elif opt in ("-d", "--dir"):
                    self.dir = arg
                elif opt in ("-r", "--recursion"):
                    self.recursion = int(arg)
                elif opt in ("-l", "--log"):
                    self.log = int(arg)
                elif opt in ("-o", "--output"):
                    self.output = arg

    def _is_image(self, filename):
        """判断是否是当前支持的图片类型"""
        _, suffix = os.path.splitext(filename)
        return str(suffix).lower() in SUFFIX

    def _is_png(self, suffix):
        return str(suffix).lower() == '.png'

    def _image_compress(self, img_item):
        if not os.path.exists(self.output):
            os.makedirs(self.output)
        
        img_item_path = Path(img_item)

        img_file_name = img_item_path.name
        img_item_data = {'fileName': img_file_name, 'sourcePath': img_item}
        img: Image.Image = Image.open(img_item_path)
        shotname = '' # 文件名
        extension = '' # 扩展名
        (shotname, extension) = os.path.splitext(img_file_name)
        # 获取压缩前的文件byte
        byteSizeBefore = len(img.fp.read())
        img_item_data['byteSizeBefore'] = byteSizeBefore
        
        # 只压缩大于100KB图片
        if byteSizeBefore < 102400:
            return

        # 区别jpg、png
        if self._is_png(extension):
            if self.jpga > 0:
                img = img.convert('RGB')
                extension = ".jpg"
            else:
                img = img.quantize(colors=256)
        save_file = os.path.join(self.output, shotname + extension)
        img_item_data['targetPath'] = save_file
        img.save(save_file, quality=self.quality, optimize=True, subsampling=self.subsampling)
        byteSizeAfter = os.path.getsize(save_file)
        img_item_data['byteSizeAfter'] = byteSizeAfter
        self.images_info.append(img_item_data)

        if self.log:
            print(img_file_name, "压缩前:", str(convert_mb_kb(byteSizeBefore)), "压缩后:", str(convert_mb_kb(byteSizeAfter)))
            print("-"*70)


    def start(self):
        """开始压缩图片"""
        if self.dir:
            # 指定文件夹
            for dirpath, dirname, filenames in os.walk(self.dir):
                if Path(dirpath) != Path(self.output):
                    for filename in filenames:
                        if self._is_image(filename):
                            self._image_compress(os.path.join(dirpath, filename))
                # 是否递归遍历子文件夹
                if not self.recursion:
                    del dirname[:]
        elif self.imgs:
            # 传入单张或多张图片
            for img in self.imgs.split(','):
                self._image_compress(img)

        if self.images_info:
            all_data = {
                'output': self.output,
                'images_info': self.images_info,
                'images': [i['targetPath'] for i in self.images_info]
            }
            print(json.dumps(all_data, ensure_ascii=False))
        else:
            print('')


if __name__ == '__main__':
    ImageCompression(sys.argv[1:]).start()